#!/bin/bash
# description: Add ssh public key to ansible inventory file
# version: 0.0.1
# author: Chuck<zhangchiqian@foxmail.com>
# date: 2020-12-28

INVENTORY_FILE="/etc/ansible/hosts"
PASSWD="root"

#HOST_GROUP=Load_balance_group

cat << EOF 
+------------------------------------------+
|   Config ssh public key ...   
+------------------------------------------+
EOF

function conf_sshkey_to_ip(){
    ips="$*"  
    for i in $ips;do
        /usr/bin/ping -f -c3 $i
        if [ $? -eq 0 ];then
            case $METHOD in 
                add)
                     echo "Add pub key to $i,Please wait"
                     /usr/bin/sshpass -p "$PASSWD" /usr/bin/ssh-copy-id -i ~/.ssh/id_rsa.pub root@$i
                     if [ $? -ne 0 ];then
                         sed -ri '/^'$i'.*/d' ~/.ssh/known_hosts
                         /usr/bin/ssh-copy-id -i ~/.ssh/id_rsa.pub root@$i
                     fi
                     ;;
                delete)
                     echo "Delete all pub key from $i,Please wait"
                     /usr/bin/ssh root@$i "echo '' > ~/.ssh/authorized_keys"
                     ;;
             esac
        else
            echo "$i not reachable"
        fi
    done
}

function conf_sshkey_to_group(){
        hosts_ips=$( awk 'BEGIN{a=0}{ if($1~/^\[[[:space:]]*'$1'[[:space:]]*\]/){a=1;next} }{ if(a!=0 && $1!~/^#.*/ && $1!~/^\[.*\]/)print $1;else if($1~/^\[.*\]/)a=0 }' $INVENTORY_FILE )
        conf_sshkey_to_ip $hosts_ips
}

function conf_sshkey(){
    if [ ! -e ~/.ssh/id_rsa -a ! -e ~/.ssh/id_rsa.pub ];then
        echo "Key not exist,generate rsa key pair"
        /usr/bin/ssh-keygen -t rsa -P ""
    fi
    if  $(echo "$1"|grep  -P '^((2(5[0-5]|[0-4]\d))|[0-1]?\d{1,2})(\.((2(5[0-5]|[0-4]\d))|[0-1]?\d{1,2})){3}$' &>/dev/null );then
        conf_sshkey_to_ip $*
    elif $(grep  '^\[[[:space:]]*'$1':children[[:space:]]*\]' $INVENTORY_FILE &>/dev/null );then 
        group_names=$( awk 'BEGIN{a=0}{ if($1~/^\[[[:space:]]*'$1':children[[:space:]]*\]/){a=1;next} }{ if(a!=0 && $1!~/^#.*/ && $1!~/^\[.*\]/)print $1;else if($1~/^\[.*\]/)a=0 }' $INVENTORY_FILE )
        for group_name in $group_names;do
             conf_sshkey_to_group $group_name
        done
    elif $(grep  '^\[[[:space:]]*'$1'[[:space:]]*\]' $INVENTORY_FILE &>/dev/null );then
        conf_sshkey_to_group $1
    else
        echo "$1 not found!"
        exit 3
    fi
}

main(){
    METHOD="add"
    if [ $# -gt 2 -o $# -eq 0 ];then
        echo "          the number of  parameters error "
        main help
        exit 2
    elif [ $# -eq 1 ];then
         case $1 in
             list)
                  awk '{if($1!~/^#.*/ && $1!~/^[[:space:]]+$/ && $1!~/^$/){if($1!~/^\[.*\]/){printf "%+20s\n",$1}else print $1}}' $INVENTORY_FILE
                  ;;
             help)
                 cat << EOF
                 Usage:
                     list    list Group Name and Hosts ips
                     [Option] <Group Name>    add ssh public key to hosts in <Group Name>
                     [Option] <IP Address>    add ssh public key to ip address
                     help    See command help
                 Option:
                     -d Delete the ssh public key on the choosen hosts.    
EOF
                 exit 0
                 ;;
             *)
                 conf_sshkey $*
                 ;;
         esac
    elif [ $# -eq 2 ];then
        case $1 in
            -d)
                METHOD="delete"
                conf_sshkey $2
                ;;
             *)
                echo "Wrong Option,abort excuting!" && main help
                exit 2
                ;;
        esac
    fi
}
main $* 
